#!/usr/bin/env python3

"""
Lerobot Control Launch - 启动 lerobot 硬件接口和相关节点
"""

from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
from launch import conditions
import os


def generate_launch_description():
    # Get package directories
    lerobot_driver_share = get_package_share_directory("scepter_lerobot_driver")
    moveit_config_share = get_package_share_directory("scepter_moveit_config")
    
    # Launch arguments
    robot_ip_arg = DeclareLaunchArgument(
        "robot_ip",
        default_value="localhost",
        description="IP address of the robot"
    )
    
    config_file_arg = DeclareLaunchArgument(
        "config_file",
        default_value="",
        description="Lerobot config file path"
    )
    
    update_rate_arg = DeclareLaunchArgument(
        "update_rate",
        default_value="100.0",
        description="Update rate for lerobot bridge (Hz)"
    )
    
    use_rviz_arg = DeclareLaunchArgument(
        "use_rviz",
        default_value="true",
        description="Whether to start RViz"
    )
    
    # Get launch configurations
    robot_ip = LaunchConfiguration("robot_ip")
    config_file = LaunchConfiguration("config_file")
    update_rate = LaunchConfiguration("update_rate")
    use_rviz = LaunchConfiguration("use_rviz")
    
    # ROS2 controllers configuration
    ros2_controllers_path = os.path.join(
        moveit_config_share, "config", "ros2_controllers.yaml"
    )
    
    # Robot description for lerobot
    robot_description_file = os.path.join(
        moveit_config_share, "config", "lerobot.arm.urdf.xacro"
    )
    
    # 1. Lerobot Bridge Node
    lerobot_bridge_node = Node(
        package="scepter_lerobot_driver",
        executable="lerobot_bridge_node.py",
        name="lerobot_bridge_node",
        output="screen",
        parameters=[
            {"robot_config": "so101"},
            {"update_rate": update_rate},
            {"command_timeout": 1.0},
        ],
    )
    
    # 2. Robot State Publisher
    robot_state_publisher = Node(
        package="robot_state_publisher",
        executable="robot_state_publisher",
        name="robot_state_publisher",
        output="screen",
        parameters=[
            {
                "robot_description": open(robot_description_file).read().format(
                    robot_ip=robot_ip,
                    config_file=config_file
                )
            }
        ],
    )
    
    # 3. Controller Manager
    controller_manager = Node(
        package="controller_manager",
        executable="ros2_control_node",
        parameters=[
            {"robot_description": open(robot_description_file).read()},
            ros2_controllers_path,
        ],
        output="screen",
    )
    
    # 4. Joint State Broadcaster Spawner
    joint_state_broadcaster_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["joint_state_broadcaster", "--controller-manager", "/controller_manager"],
        output="screen",
    )
    
    # 5. Arm Controller Spawner
    arm_controller_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["arm_group_controller", "-c", "/controller_manager"],
        output="screen",
    )
    
    # 6. Gripper Controller Spawner
    gripper_controller_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["gripper_group_controller", "-c", "/controller_manager"],
        output="screen",
    )
    
    # 7. Static TF (world to base_link)
    static_tf = Node(
        package="tf2_ros",
        executable="static_transform_publisher",
        name="static_transform_publisher",
        output="log",
        arguments=["--frame-id", "world", "--child-frame-id", "base_link"],
    )
    
    # 8. RViz (optional)
    rviz_config_file = PathJoinSubstitution([
        FindPackageShare("scepter_moveit_config"),
        "config",
        "moveit.rviz"
    ])
    
    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        arguments=["-d", rviz_config_file],
        condition=conditions.IfCondition(use_rviz),
    )
    
    return LaunchDescription([
        # Launch arguments
        robot_ip_arg,
        config_file_arg,
        update_rate_arg,
        use_rviz_arg,
        
        # Nodes
        lerobot_bridge_node,
        robot_state_publisher,
        static_tf,
        controller_manager,
        joint_state_broadcaster_spawner,
        arm_controller_spawner,
        gripper_controller_spawner,
        rviz_node,
    ])